#include <complex>
#include <iostream>
#include <vector>
#include <random>
#include <mc_runtime.h>
#include "mcfft.h"

#define NX 4
#define NY 4
#define NRANK 2
#define BATCH 1

using namespace std;
int main(void)
{
    std::cout<<"mcFFT real/complex 2d FFT example\n";
    mcfftHandle plan;
    mcError_t err;
    int n[NRANK] = {NX, NY};
    // Initialize data on the host:
    std::vector<float>data(NX*NY*BATCH);
    std::vector<float>outData(NX*NY*BATCH);
    std::cout<<"intput array: \n";
    for(int b=0;b<BATCH;b++)
    {
        for(int n=0;n<NX;n++)
        {
            for(int m=0;m<NY;m++)
            {
                const size_t pos = (b * NX + n) * NY + m;
                const size_t genValue = n*NY+m;
                std::mt19937 gen(genValue);
                const float  val = (float)gen() / (float)gen.max();
                data[pos]=val;
                std::cout<<data[pos]<<" ";
            } 
        }
    }
    // Create device object and plan:
    mcfftReal *devPtrData=nullptr;
    err=mcMalloc((void**)&devPtrData, sizeof(mcfftReal)*NX*NY*BATCH);
    if (err != mcSuccess)
    {
	    fprintf(stderr,"Error: Failed to allocate\n");
	    return EXIT_FAILURE;	
    }
    mcfftComplex *devPtrOutData=nullptr;
    err=mcMalloc((void**)&devPtrOutData, sizeof(mcfftComplex)*NX*(NY/2+1)*BATCH);
    if (err != mcSuccess)
    {
	    fprintf(stderr,"Error: Failed to allocate\n");
	    return EXIT_FAILURE;	
    }
    // Create a 2D FFT plan:
    if (mcfftPlanMany(&plan, NRANK, n,
				        NULL, 1, 0, // *inembed, istride, idist 
				        NULL, 1, 0, // *onembed, ostride, odist
    				    MCFFT_R2C,BATCH) != MCFFT_SUCCESS){
    	   fprintf(stderr, "MCFFT Error: Unable to create plan\n");
    	   return EXIT_FAILURE;	
    }
    //copy data:
     err=mcMemcpy(devPtrData, data.data(), sizeof(mcfftReal)*NX*NY*BATCH, mcMemcpyHostToDevice);
    if (err != mcSuccess)
    {
	    fprintf(stderr,"Error: Failed to copy host to device\n");
	    return EXIT_FAILURE;	
    }
    /* Note:
    *  Different pointers to input and output arrays implies out-place transformation
    */
   //forward tranform R2C:
    if (mcfftExecR2C(plan, devPtrData, devPtrOutData) != MCFFT_SUCCESS)
    {
	    fprintf(stderr,"MCFFT error: ExecC2C Forward failed\n");
	    return EXIT_FAILURE;	
    }
    // Create a 2D FFT plan:
    if (mcfftPlanMany(&plan, NRANK, n,
				        NULL, 1, 0, // *inembed, istride, idist 
				        NULL, 1, 0, // *onembed, ostride, odist
    				    MCFFT_C2R,BATCH) != MCFFT_SUCCESS){
    	fprintf(stderr, "MCFFT Error: Unable to create plan\n");
    	return EXIT_FAILURE;	
    }
   //inverse transform C2R:
    if (mcfftExecC2R(plan, devPtrOutData, devPtrData) != MCFFT_SUCCESS)
    {
	    fprintf(stderr,"MCFFT error: ExecC2C Inverse failed\n");
	    return EXIT_FAILURE;	
    }
    /* 
    *  Results may not be immediately available so block device until all 
    *  tasks have completed
    */
    if (mcDeviceSynchronize() != mcSuccess){
	    fprintf(stderr,"Error: Failed to synchronize\n");
	    return EXIT_FAILURE;	
    }
    err=mcMemcpy(outData.data(), devPtrData, sizeof(mcfftReal)*NX*NY*BATCH,mcMemcpyDeviceToHost);
    if (err  != mcSuccess){
	    fprintf(stderr,"Error: Failed to copy device to host\n");
	    return EXIT_FAILURE;	
    }


    const float overN = 1.0/(NX*NY);
    float       error = 0;
    for(size_t i = 0; i < data.size(); i++)
    {
       double diff = std::fabs(outData[i]*overN - data[i]);
       if(diff > error)
        {
            error = diff;
        }
    }
    std::cout<<"\ntranformed back: \n";
    for(size_t i = 0; i < outData.size(); i++)
    {
        std::cout<<outData[i]*overN<<" ";
    }
    std::cout<<"\nMaximum error: "<<error<<std::endl;
    mcfftDestroy(plan);
    mcFree(devPtrData);
    mcFree(devPtrOutData);
}

